import os
import pandas as pd
import sqlite3
import sqlalchemy

_csv_dir = './marketData/csv/'
_marketDataDBDir = './marketData/'
_marketDataDBName = 'marketData.db'
_marketDataDBPath = _marketDataDBDir + _marketDataDBName

_marketDataDBNameJQ = 'marketDataJQ.db'
_marketDataDBPathJQ = _marketDataDBDir + _marketDataDBNameJQ

_marketDataDBDownloaded = 'C:/Users/yang/Downloads/' + _marketDataDBNameJQ

# 价-量数据的 csv 标准
_price_volume_csv_columns = ['time', 'open', 'high', 'low', 'close', 'volume']
_price_volume_columns_withoutTime = ['open', 'high', 'low', 'close', 'volume']

# 价-量数据的 dataFrame 标准
_price_volume_index_column = 'time'
_price_volume_columns = [
    'code', 'frequency', 'open', 'high', 'low', 'close', 'volume']
_price_volume_columns_dtypes = {'open': 'float32',
                                'high': 'float32',
                                'low': 'float32',
                                'close': 'float32',
                                'volume': 'float64'}
_price_volume_columns_round = {'open': 2,
                               'high': 2,
                               'low': 2,
                               'close': 2,
                               'volume': 0}


# 价-量数据标准 csv 读写
def _to_price_volume_csv(time_index,
                         open_col,
                         high_col,
                         low_col,
                         close_col,
                         volume_col,
                         code,
                         freq,
                         output_dir=_csv_dir):
    '''
    time_index, open_col, high_col, low_col, close_col, volume_col
    需要长度相等的list数据
    code, freq 用于为文件命名

    return: output_path
    '''

    if type(time_index) == list and type(open_col) == list and \
            type(high_col) == list and type(low_col) == list and \
            type(close_col) == list and type(volume_col) == list and \
            len(time_index) == len(open_col) and \
            len(time_index) == len(high_col) and \
            len(time_index) == len(low_col) and \
            len(time_index) == len(close_col) and \
            len(time_index) == len(volume_col) and len(time_index):

        col_data = {
            'open': open_col,
            'high': high_col,
            'low': low_col,
            'close': close_col,
            'volume': volume_col
        }
        pv_df = pd.DataFrame(data=col_data, index=time_index)
        pv_df.index.name = 'time'

        if os.path.isdir(output_dir) is False:
            os.makedirs(output_dir)
        output_filename = code + '_' + freq + '.csv'
        output_path = output_dir + output_filename
        pv_df.to_csv(output_path)
        return output_path
    else:
        raise Exception('input not satisfied')


def _price_volume_from_csv(csv_path='', code='', frequency=''):
    if csv_path == '':
        raise Exception('none csv_path')
    else:
        print('price_volume data loading: ' + csv_path)

    # 读入
    pv_df = pd.read_csv(
        csv_path,
        index_col=_price_volume_index_column,
        usecols=_price_volume_csv_columns,
        parse_dates=[_price_volume_index_column],
        infer_datetime_format=True,
        dtype=_price_volume_columns_dtypes)[_price_volume_columns_withoutTime]

    if pv_df.empty is True:
        raise Exception('empty csv')

    pv_df['code'] = code
    pv_df['frequency'] = frequency

    # 最后处理
    pv_df = pv_df[_price_volume_columns]
    pv_df = pv_df.astype(_price_volume_columns_dtypes)
    pv_df = pv_df.round(_price_volume_columns_round)

    return pv_df


# 价-量数据标准 SQLite 读写
def _tableName(code: str, frequency: str) -> str:
    return '%s_%s' % (code, frequency)


def _fromTableName(tableName: str) -> tuple[str, str]:
    '''
    :param tableName:
    :return: code, frequency
    '''
    p = 0
    for char in tableName:
        p += 1
        if char == '_':
            return tableName[:p - 1], tableName[p:]


def _pvToSQLiteDB(pv: pd.DataFrame, DBPath: str) -> None:
    if pv.size < 1:
        raise Exception('empty pv.')
    DBPath = os.path.abspath(DBPath)
    if not os.path.exists(DBPath):
        dir = os.path.split(DBPath)[0]
        if not os.path.exists(dir):
            os.makedirs(dir)
        sqlite3.connect(DBPath).commit()

    tableName = _tableName(pv['code'][0], pv['frequency'][0])

    pv = pv[_price_volume_columns_withoutTime]
    engine = sqlalchemy.create_engine('sqlite:///' + DBPath)
    tableNames = engine.table_names()
    if tableName in tableNames:
        print('price-volume data overwriting: %s --> %s' % (tableName, DBPath))
        with engine.connect() as conn:
            with conn.begin():
                sqlText = 'DROP TABLE IF EXISTS "%s"' % tableName
                conn.execute(sqlText)
                pv.to_sql(tableName, engine)
    else:
        print('price-volume data saving: %s --> %s' % (tableName, DBPath))
        with engine.connect() as conn:
            with conn.begin():
                pv.to_sql(tableName, engine)

    return None


def _pvFromSQLiteDB(code: str,
                    frequency: str,
                    DBPath: str = _marketDataDBPath,
                    start: str = '',
                    end: str = '',
                    bars: int = 0) -> pd.DataFrame:
    DBPath = os.path.abspath(DBPath)
    if not os.path.exists(DBPath):
        raise Exception('DBPath not exists: %s' % DBPath)

    tableName = _tableName(code, frequency)
    pv = None

    # print('price_volume data loading: ' + tableName)
    engine = sqlalchemy.create_engine('sqlite:///' + DBPath)
    time_column = _price_volume_index_column
    with engine.connect() as conn:
        with conn.begin():
            if not bars:
                if not start and not end:
                    pv = pd.read_sql_table(tableName, conn)
                elif not start:
                    sql = ('SELECT * FROM "%s"' +
                           'WHERE datetime("%s") <= datetime("%s")') % (
                        tableName, time_column,
                        pd.Timestamp(end).strftime('%Y-%m-%d %H:%M:%S'))
                    pv = pd.read_sql_query(sql, conn)
                elif not end:
                    sql = ('SELECT * FROM "%s"' +
                           'WHERE datetime("%s") >= datetime("%s")') % (
                        tableName, time_column,
                        pd.Timestamp(start).strftime('%Y-%m-%d %H:%M:%S'))
                    pv = pd.read_sql_query(sql, conn)
                else:
                    sql = ('SELECT * FROM "%s"' +
                           'WHERE datetime("%s") >= datetime("%s")' +
                           'AND datetime("%s") <= datetime("%s")') % (
                        tableName, time_column,
                        pd.Timestamp(start).strftime('%Y-%m-%d %H:%M:%S'),
                        time_column,
                        pd.Timestamp(end).strftime('%Y-%m-%d %H:%M:%S'))
                    pv = pd.read_sql_query(sql, conn)
            else:
                sql = ('SELECT * FROM "%s" ' +
                       'WHERE datetime("%s") <= datetime("%s") ' +
                       'ORDER BY datetime("%s") DESC LIMIT %d') % (
                    tableName, time_column,
                    pd.Timestamp(end).strftime('%Y-%m-%d %H:%M:%S'),
                    time_column, bars)
                pv = pd.read_sql_query(sql, conn)
                if len(pv) != bars:
                    raise Exception('no enough data')
    pv[time_column] = pd.to_datetime(pv[time_column])
    pv = pv.set_index(time_column)
    pv = pv.sort_index(ascending=True)
    pv['code'] = code
    pv['frequency'] = frequency

    pv = pv[_price_volume_columns]
    pv = pv.astype(_price_volume_columns_dtypes)
    pv = pv.round(_price_volume_columns_round)

    return pv


def transferPVTable(fromDBPath: str = _marketDataDBDownloaded,
                    toDBPath: str = _marketDataDBPath) -> None:
    fromDBPath = os.path.abspath(fromDBPath)
    toDBPath = os.path.abspath(toDBPath)
    if not os.path.exists(fromDBPath):
        raise Exception('DBPath not exists: %s' % fromDBPath)
    if not os.path.exists(toDBPath):
        dir = os.path.split(toDBPath)[0]
        if not os.path.exists(dir):
            os.makedirs(dir)
        sqlite3.connect(toDBPath).commit()

    engine = sqlalchemy.create_engine('sqlite:///' + fromDBPath)
    tableNames = engine.table_names()
    if len(tableNames) == 0:
        print('no tables in: %s' % fromDBPath)
        return None
    else:
        for tableName in tableNames:
            code, freq = _fromTableName(tableName)
            print('transferring table: %s --> %s --> %s' %
                  (fromDBPath, tableName, toDBPath))
            _pvToSQLiteDB(_pvFromSQLiteDB(code, freq, fromDBPath), toDBPath)
        os.remove(fromDBPath)
        print('transferring complete.')
